package com.me.spider;

import javax.crypto.Cipher;
import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.security.GeneralSecurityException;
import java.security.MessageDigest;
import java.security.spec.KeySpec;

/**
 * 解密微信本地文件
 */
public class WechatDecryptor {

    // 定义基础加密参数和设置
    private static final int DEFAULT_ITER = 1000;
    private static final int KEY_SIZE = 256;
    private static final int DEFAULT_PAGE_SIZE = 4096;
    private static final int IV_SIZE = 16;
    private static final int HMAC_SHA1_SIZE = 20;
    private static final String SQLITE_FILE_HEADER = "SQLite format 3\0";

    // 定义加密用的密码
    private static final byte[] PASS = {
            (byte) 0x53, (byte) 0xE9, (byte) 0xBF, (byte) 0xB2, (byte) 0x3B, (byte) 0x72, (byte) 0x41, (byte) 0x95,
            (byte) 0xA2, (byte) 0xBC, (byte) 0x6E, (byte) 0xB5, (byte) 0xBF, (byte) 0xEB, (byte) 0x06, (byte) 0x10,
            (byte) 0xDC, (byte) 0x21, (byte) 0x64, (byte) 0x75, (byte) 0x6B, (byte) 0x9B, (byte) 0x42, (byte) 0x79,
            (byte) 0xBA, (byte) 0x32, (byte) 0x15, (byte) 0x76, (byte) 0x39, (byte) 0xA4, (byte) 0x0B, (byte) 0xB1
    };

    public static void main(String[] args) {
        // 检查并获取文件名参数，或提示用户输入文件名
        if (args.length < 1) {
            System.out.println("请输入文件名:");
            BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));
            try {
                String dbFileName = reader.readLine();
                decryptDb(dbFileName);
            } catch (IOException e) {
                e.printStackTrace();
            }
        } else {
            decryptDb(args[0]);
        }
    }

    /**
     * 解密数据库文件的主方法
     *
     * @param dbFileName 要解密的数据库文件名
     */
    private static void decryptDb(String dbFileName) {
        try (FileInputStream fileInputStream = new FileInputStream(dbFileName);
             FileOutputStream fileOutputStream = new FileOutputStream("dec_" + dbFileName, true)) {

            // 从文件中读取盐值
            byte[] salt = new byte[IV_SIZE];
            if (fileInputStream.read(salt) != IV_SIZE) {
                throw new IOException("Failed to read the salt from the file.");
            }

            // 根据密码和盐值生成AES密钥
            SecretKeySpec keySpec = new SecretKeySpec(generateKey(PASS, salt), "AES");

            byte[] pageBuffer = new byte[DEFAULT_PAGE_SIZE];
            int pageNumber = 1;

            // 循环读取文件的每一页并进行解密处理
            while (fileInputStream.read(pageBuffer) == DEFAULT_PAGE_SIZE) {
                System.out.println("解密数据页: " + pageNumber);

                // 对数据页面进行HMAC验证
                if (!verifyHMAC(pageBuffer, pageNumber, keySpec)) {
                    System.err.println("HMAC verification failed for page: " + pageNumber);
                    return;
                }

                // 提取IV用于AES解密
                byte[] iv = extractIV(pageBuffer);

                // 初始化AES解密Cipher
                Cipher cipher = Cipher.getInstance("AES/CBC/NoPadding");
                cipher.init(Cipher.DECRYPT_MODE, keySpec, new IvParameterSpec(iv));

                // 执行解密操作
                byte[] decryptedPage = cipher.doFinal(pageBuffer, IV_SIZE, DEFAULT_PAGE_SIZE - IV_SIZE - HMAC_SHA1_SIZE);

                // 处理SQLite文件头（只有第一页）
                if (pageNumber == 1) {
                    System.arraycopy(SQLITE_FILE_HEADER.getBytes(), 0, decryptedPage, 0, SQLITE_FILE_HEADER.length());
                }

                // 将解密后的数据写入输出文件
                fileOutputStream.write(decryptedPage);
                pageNumber++;
            }

            System.out.println("\n解密成功!");

        } catch (IOException | GeneralSecurityException e) {
            e.printStackTrace();
        }
    }

    /**
     * "53 E9 BF B2 3B 72 41 95 A2 BC 6E B5 BF EB 06 10 DC 21 64 75 6B 9B 42 79 BA 32 15 76 39 A4 0B B1"
     * 把十六进制字符串文本转换成 byte[]
     */
    public static byte[] convertHexStringToByteArray(String hexString) {
        // 移除字符串中的空格
        hexString = hexString.replaceAll("\\s", "");

        // 初始化byte数组，长度为字符串长度的一半（因为每两个字符表示一个byte）
        byte[] byteArray = new byte[hexString.length() / 2];

        // 循环处理字符串，每次处理两个字符
        for (int i = 0; i < byteArray.length; i++) {
            int index = i * 2;
            // 将每两个字符转换为一个十六进制整数，然后转换为byte
            int val = Integer.parseInt(hexString.substring(index, index + 2), 16);
            byteArray[i] = (byte) val;
        }

        return byteArray;
    }

    /**
     * 使用 PBKDF2算法 和 HMAC-SHA1 生成 AES密钥
     *
     * @param pass 密码
     * @param salt 盐值
     * @return 生成的AES密钥
     * @throws GeneralSecurityException 加密操作异常
     */
    private static byte[] generateKey(byte[] pass, byte[] salt) throws GeneralSecurityException {
        SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
        KeySpec spec = new PBEKeySpec(new String(pass).toCharArray(), salt, DEFAULT_ITER, KEY_SIZE);
        SecretKey tmp = factory.generateSecret(spec);
        return tmp.getEncoded();
    }

    /**
     * 验证页面数据的HMAC-SHA1签名
     *
     * @param pageBuffer 读取的页面数据
     * @param pageNumber 页面编号
     * @param keySpec 密钥规范
     * @return 验证结果
     * @throws GeneralSecurityException 加密操作异常
     */
    private static boolean verifyHMAC(byte[] pageBuffer, int pageNumber, SecretKeySpec keySpec) throws GeneralSecurityException {
        Mac mac = Mac.getInstance("HmacSHA1");
        mac.init(keySpec);

        // 计算该页的HMAC-SHA1值，并与页面中的HMAC值比较
        byte[] dataToVerify = new byte[DEFAULT_PAGE_SIZE - HMAC_SHA1_SIZE];
        System.arraycopy(pageBuffer, 0, dataToVerify, 0, dataToVerify.length);
        byte[] calculatedHMAC = mac.doFinal(dataToVerify);

        byte[] hmacFromPage = new byte[HMAC_SHA1_SIZE];
        System.arraycopy(pageBuffer, DEFAULT_PAGE_SIZE - HMAC_SHA1_SIZE, hmacFromPage, 0, HMAC_SHA1_SIZE);

        // 使用安全的方法比较两个字节数组的内容
        return MessageDigest.isEqual(calculatedHMAC, hmacFromPage);
    }

    /**
     * 从页面数据中提取IV（初始化向量）
     *
     * @param pageBuffer 页面数据
     * @return 提取的IV
     */
    private static byte[] extractIV(byte[] pageBuffer) {
        byte[] iv = new byte[IV_SIZE];
        System.arraycopy(pageBuffer, DEFAULT_PAGE_SIZE - IV_SIZE - HMAC_SHA1_SIZE, iv, 0, IV_SIZE);
        return iv;
    }

}
